package uricontrol

import (
	"errors"
	"fmt"
	"gitee.com/baal1990/gitee.com/baal1990/pkg/cache"
	"github.com/gomodule/redigo/redis"
	"github.com/mailru/easyjson"
	"time"
)

type UriLimiter struct {
	Ref      int   `json:"ref"`
	UnixNano int64 `json:"unix_nano"`
}

// UriControl 安全校验，限制终端ip，防止暴力破解、短信攻击
func UriControl(expire, count int, clientIp, uri string) (err error) {
	if clientIp == "" || uri == "" || expire <= 0 || count <= 0 {
		return errors.New(fmt.Sprintf("invalid control with expire[%d], count[%d], clientIp[%s], uri[%s]", expire, count, clientIp, uri))
	}
	key := RedisUriLimiter + clientIp
	var limiter UriLimiter
	limiter, err = GetUriLimiter(key, uri)
	if err != nil {
		return
	}
	// 是否过期
	now := time.Now().UnixNano()
	if now <= limiter.UnixNano { // 未过期
		if limiter.Ref >= count { // 超限
			return errors.New(fmt.Sprintf("too many request[%s]", uri))
		} else { //未超限
			limiter.Ref++
			if err = SetUriLimiter(limiter, key, uri); err != nil {
				return
			}
		}
	} else { // 过期
		limiter.Ref = 1
		limiter.UnixNano = time.Now().Add(time.Second * time.Duration(expire)).UnixNano()
		if err = SetUriLimiter(limiter, key, uri); err != nil {
			return
		}
	}
	return nil
}

// GetUriLimiter 读取UriLimiter
func GetUriLimiter(key, field string) (limiter UriLimiter, err error) {
	var jsonBuf []byte
	var reply interface{}
	if reply, err = cache.Do("HGET", key, field); err != nil {
		return limiter, errors.New(fmt.Sprintf("HGET [%s].[%s] fail, [%v]", key, field, err))
	}
	if jsonBuf, err = redis.Bytes(reply, nil); err == nil {
		if err = easyjson.Unmarshal(jsonBuf, &limiter); err != nil {
			return limiter, errors.New(fmt.Sprintf("easyjson.Unmarshal fail, [%v]", err))
		}
	}
	return limiter, nil
}

// SetUriLimiter UriLimiter存入缓存
func SetUriLimiter(limiter UriLimiter, key, field string) (err error) {
	var jsonBuf []byte
	if jsonBuf, err = easyjson.Marshal(limiter); err != nil {
		return errors.New("easyjson.Marshal fail, " + err.Error())
	}
	if _, err = cache.Do("HSET", key, field, jsonBuf); err != nil {
		return errors.New("HSET fail fail, " + err.Error())
	}
	return nil
}
